from typing import Dict, Any, List, Optional
import chromadb
from chromadb.config import Settings
from agents.components.base import BaseComponent, ComponentConfig, ComponentType
from agents.components.embed.base_embedding_component import EmbeddingComponent

class ChromaComponent(BaseComponent):
    """ChromaDB向量检索组件实现"""
    
    @classmethod
    def get_component_type(cls) -> ComponentType:
        return ComponentType.VECTOR_DB
    
    @classmethod
    def get_component_name(cls) -> str:
        return "chroma"
    
    @classmethod
    def get_component_description(cls) -> str:
        return "基于ChromaDB的向量检索组件"
    
    @classmethod
    def get_config_schema(cls) -> Dict[str, Any]:
        return {
            "type": "object",
            "properties": {
                "persist_directory": {
                    "type": "string",
                    "description": "ChromaDB持久化目录"
                },
                "collection_name": {
                    "type": "string",
                    "description": "集合名称"
                },
                "page_size": {
                    "type": "integer",
                    "description": "默认分页大小",
                    "default": 10
                },
                "n_results": {
                    "type": "integer",
                    "description": "默认返回结果数量",
                    "default": 10
                },
                "embedding_component": {
                    "type": "string",
                    "description": "使用的embedding组件ID"
                }
            },
            "required": ["collection_name", "embedding_component"]
        }
    
    @classmethod
    def get_input_schema(cls) -> Dict[str, Any]:
        return {
            "type": "object",
            "properties": {
                "query_text": {
                    "type": "string",
                    "description": "查询文本"
                },
                "page": {
                    "type": "integer",
                    "description": "页码",
                    "default": 1
                },
                "page_size": {
                    "type": "integer",
                    "description": "每页大小"
                },
                "search_type": {
                    "type": "string",
                    "description": "搜索类型",
                    "enum": ["hybrid", "vector", "text"],
                    "default": "hybrid"
                },
                "record_id": {
                    "type": "string",
                    "description": "记录ID"
                },
                "metadata_filter": {
                    "type": "object",
                    "description": "元数据过滤条件"
                },
                "document_filter": {
                    "type": "object",
                    "description": "文档过滤条件"
                },
                "n_results": {
                    "type": "integer",
                    "description": "返回结果数量"
                }
            },
            "required": ["query_text"]
        }
    
    @classmethod
    def get_output_schema(cls) -> Dict[str, Any]:
        return {
            "type": "object",
            "properties": {
                "records": {
                    "type": "array",
                    "items": {
                        "type": "object",
                        "properties": {
                            "id": {"type": "string"},
                            "text": {"type": "string"},
                            "metadata": {"type": "object"},
                            "distance": {"type": "number"}
                        }
                    }
                },
                "total": {"type": "integer"},
                "message": {"type": "string"}
            },
            "required": ["records", "total"]
        }
    
    def __init__(self, config: ComponentConfig):
        super().__init__(config)
        self.persist_directory = config.config.get("persist_directory", "./chroma_db")
        self.collection_name = config.config.get("collection_name", "")
        self.default_page_size = config.config.get("page_size", 10)
        self.default_n_results = config.config.get("n_results", 10)
        self.embedding_component_id = config.config.get("embedding_component")
        
        # 初始化ChromaDB客户端
        self.client = chromadb.Client(Settings(
            persist_directory=self.persist_directory,
            is_persistent=True
        ))
        
        # 获取或创建集合
        self.collection = self.client.get_or_create_collection(
            name=self.collection_name
        )
        
    def set_embedding_component(self, component: EmbeddingComponent):
        """设置embedding组件"""
        self.embedding_component = component
        
    async def execute(self, input_data: Dict[str, Any]) -> Dict[str, Any]:
        """执行向量检索"""
        try:
            # 获取检索参数
            query_text = input_data.get("query_text", "")
            page = input_data.get("page", 1)
            page_size = input_data.get("page_size", self.default_page_size)
            search_type = input_data.get("search_type", "hybrid")
            record_id = input_data.get("record_id")
            metadata_filter = input_data.get("metadata_filter")
            document_filter = input_data.get("document_filter")
            n_results = input_data.get("n_results", self.default_n_results)
            
            # 执行向量检索
            if record_id:
                # 按ID查询
                results = self.collection.get(
                    ids=[record_id],
                    where=metadata_filter
                )
            else:
                # 使用embedding组件获取查询文本的向量
                embedding_result = await self.embedding_component.execute({
                    "texts": [query_text]
                })
                query_embedding = embedding_result["embeddings"][0]
                
                # 按向量查询
                results = self.collection.query(
                    query_embeddings=[query_embedding],
                    n_results=n_results,
                    where=metadata_filter,
                    where_document=document_filter
                )
            
            # 处理检索结果
            if not results or not results.get('ids'):
                return {
                    "records": [],
                    "total": 0,
                    "message": "未找到匹配的记录"
                }
            
            # 格式化结果
            records = []
            for i in range(len(results['ids'])):
                records.append({
                    "id": results['ids'][i],
                    "text": results['documents'][i],
                    "metadata": results['metadatas'][i],
                    "distance": results['distances'][i] if 'distances' in results else None
                })
            
            # 分页处理
            start_idx = (page - 1) * page_size
            end_idx = start_idx + page_size
            paginated_records = records[start_idx:end_idx]
            
            return {
                "records": paginated_records,
                "total": len(records),
                "message": "检索成功"
            }
            
        except Exception as e:
            return {
                "error": str(e),
                "records": [],
                "total": 0
            }
    
    def validate_input(self, input_data: Dict[str, Any]) -> bool:
        """验证输入数据"""
        required_fields = self.config.input_schema.get("required", [])
        return all(field in input_data for field in required_fields)
    
    def validate_output(self, output_data: Dict[str, Any]) -> bool:
        """验证输出数据"""
        return "records" in output_data and isinstance(output_data["records"], list) 